from openai import OpenAI
import openai
from dotenv import load_dotenv
from app.utils.format_message import format_user_message
import tiktoken
import os
import aiohttp
import json
from typing import AsyncGenerator, Literal
import requests
import time
import traceback
import random
import math
import asyncio
from typing import Optional, Dict, List, Union, Any, Tuple

load_dotenv()


class OpenAIo3Service:
    def __init__(self):
        self.default_client = OpenAI(
            api_key=os.getenv("OPENAI_API_KEY"),
            timeout=60.0,  # 设置默认超时时间为60秒
        )
        try:
            self.encoding = tiktoken.get_encoding("o200k_base")  # Encoder for OpenAI models
            print("成功加载o200k_base编码器")
        except (requests.exceptions.SSLError, requests.exceptions.ConnectionError, 
                requests.exceptions.Timeout, Exception) as e:
            print(f"无法加载o200k_base编码器: {str(e)}，将使用cl100k_base作为备选")
            try:
                # 尝试使用cl100k_base作为备选
                self.encoding = tiktoken.get_encoding("cl100k_base")
                print("成功加载cl100k_base备选编码器")
            except Exception as e2:
                print(f"备选编码器加载失败: {str(e2)}，将使用简单令牌计数")
                # 如果还失败，就使用一个简单的基于字符的计数方法
                self.encoding = None
                print("使用简单字符计数作为编码器")
        
        self.base_url = "https://api.openai.com/v1/chat/completions"
        print(f"OpenAIo3Service 初始化完成，使用API密钥: {'已设置' if os.getenv('OPENAI_API_KEY') else '未设置'}")

    def call_o3_api(self, prompt: str, user_id: str = "", max_tokens: int = 1024, temperature: float = 0.7) -> str:
        """
        调用OpenAI的o3-mini模型API进行文本生成
        """
        print(f"开始调用o3-mini模型, prompt长度: {len(prompt)}")
        
        try:
            # 使用客户端发起请求
            response = self.default_client.chat.completions.create(
                model="o3",
                messages=[
                    {"role": "system", "content": "你是一个专业、高效、实用的AI助手。"},
                    {"role": "user", "content": prompt}
                ],
                max_tokens=max_tokens,
                n=1,
                temperature=temperature,
            )
            
            if not response.choices:
                print("API响应中没有choices")
                return ""
                
            # 返回生成的文本
            result = response.choices[0].message.content or ""
            input_tokens = self.count_tokens(prompt)
            output_tokens = self.count_tokens(result)
            print(f"o3-mini响应成功, 输入tokens: {input_tokens}, 输出tokens: {output_tokens}")
            
            return result
        except openai.APITimeoutError as e:
            print(f"OpenAI API超时: {e}")
            return "很抱歉，服务器处理超时，请稍后再试或尝试简化您的问题。"
        except requests.exceptions.Timeout as e:
            print(f"请求超时: {e}")
            return "很抱歉，服务器处理超时，请稍后再试或尝试简化您的问题。"
        except requests.exceptions.ConnectionError as e:
            print(f"连接错误: {e}")
            return "很抱歉，网络连接发生错误，请检查您的网络连接或稍后再试。"
        except openai.RateLimitError:
            print("o3-mini API达到速率限制")
            return "抱歉，服务暂时达到请求上限，请稍后再试。"
        except openai.APIError as e:
            print(f"o3-mini API错误: {str(e)}")
            return "抱歉，服务器处理请求时出现问题，请稍后再试。"
        except Exception as e:
            print(f"调用API错误: {e}")
            return f"抱歉，服务出现了问题: {str(e)}"

    async def call_o3_api_stream(
        self, messages: List[Dict], system_message: str = "", user_name: str = "", to_name: str = ""
    ):
        """
        Makes a streaming API call to OpenAI o3-mini and yields the responses.

        Args:
            messages (List[Dict]): List of messages to format into the user message
            system_message (str): The instruction/system prompt
            user_name (str): The user's name
            to_name (str): The recipient's name

        Yields:
            str: Chunks of o3-mini's response text
        """
        # Create the user message with the data
        user_message = format_user_message(messages)
        token_count = self.count_tokens(user_message) + self.count_tokens(system_message)
        print(f"[OpenAI o3] 流式API调用: 总令牌数={token_count}")

        headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {self.default_client.api_key}",
        }

        payload = {
            "model": "o3-mini",
            "messages": [
                {"role": "system", "content": system_message},
                {"role": "user", "content": user_message},
            ],
            "max_completion_tokens": 12000,
            "stream": True,
        }

        try:
            print(f"[OpenAI o3] 使用默认API密钥发送流式请求...")
            
            async with aiohttp.ClientSession() as session:
                async with session.post(
                    self.base_url, headers=headers, json=payload
                ) as response:

                    if response.status != 200:
                        error_text = await response.text()
                        print(f"[OpenAI o3] 错误响应: {error_text}")
                        raise ValueError(
                            f"OpenAI API returned status code {response.status}: {error_text}"
                        )

                    line_count = 0
                    chunk_count = 0
                    total_content = ""
                    
                    print(f"[OpenAI o3] 开始接收流式响应...")
                    
                    async for line in response.content:
                        line = line.decode("utf-8").strip()
                        if not line:
                            continue

                        line_count += 1

                        if line.startswith("data: "):
                            if line == "data: [DONE]":
                                break
                            try:
                                data = json.loads(line[6:])
                                content = (
                                    data.get("choices", [{}])[0]
                                    .get("delta", {})
                                    .get("content")
                                )
                                if content:
                                    chunk_count += 1
                                    total_content += content
                                    if chunk_count % 20 == 0:  # 每20个块打印一次进度
                                        print(f"[OpenAI o3] 已接收 {chunk_count} 个内容块")
                                    yield content
                            except json.JSONDecodeError as e:
                                print(f"[OpenAI o3] JSON解析错误: {e}, 行: {line}")
                                continue

                    response_tokens = self.count_tokens(total_content)
                    print(f"[OpenAI o3] 流式响应完成: 共 {chunk_count} 个内容块, {line_count} 行数据, 响应令牌数={response_tokens}")
                    
                    if line_count == 0:
                        print("[OpenAI o3] 警告: 流式响应未收到任何行")

        except openai.APITimeoutError as e:
            print(f"OpenAI API流式调用超时: {e}")
            yield "很抱歉，服务器处理超时，请稍后再试或尝试简化您的问题。"
        except requests.exceptions.Timeout as e:
            print(f"流式请求超时: {e}")
            yield "很抱歉，服务器处理超时，请稍后再试或尝试简化您的问题。"
        except requests.exceptions.ConnectionError as e:
            print(f"流式连接错误: {e}")
            yield "很抱歉，网络连接发生错误，请检查您的网络连接或稍后再试。"
        except openai.RateLimitError:
            print("o3-mini API达到速率限制")
            yield "抱歉，服务暂时达到请求上限，请稍后再试。"
        except openai.APIError as e:
            print(f"o3-mini API流式调用错误: {str(e)}")
            yield "抱歉，服务器处理请求时出现问题，请稍后再试。"
        except Exception as e:
            print(f"流式调用API错误: {e}")
            yield f"抱歉，服务出现了问题: {str(e)}"

    def count_tokens(self, text: str) -> int:
        """计算输入文本的令牌数量"""
        try:
            if self.encoding:
                return len(self.encoding.encode(text))
            else:
                # 如果没有可用的编码器，使用简单的字符估算法
                # 大约每4个字符等于1个token (这只是一个粗略估计)
                token_estimate = len(text) // 4
                return max(1, token_estimate)  # 确保至少返回1个token
        except Exception as e:
            print(f"计算令牌时出错: {str(e)}")
            traceback.print_exc()  # 打印完整的堆栈跟踪以便调试
            # 出错时也返回一个安全的估计值
            return max(1, len(text) // 4)
